import sqlite3
import os

def add_stream_enabled_column():
    # 获取数据库路径
    db_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'app', 'database', 'app.db')
    
    # 连接数据库
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 检查表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='ai_models'")
    if cursor.fetchone():
        # 检查字段是否已存在
        cursor.execute("PRAGMA table_info(ai_models)")
        columns = [column[1] for column in cursor.fetchall()]
        
        if 'stream_enabled' not in columns:
            print("添加 stream_enabled 字段到 ai_models 表...")
            
            # SQLite 不支持直接使用 ALTER TABLE ADD COLUMN 添加带有 NOT NULL 约束的列
            # 除非指定默认值
            cursor.execute("ALTER TABLE ai_models ADD COLUMN stream_enabled INTEGER NOT NULL DEFAULT 1")
            
            # 提交更改
            conn.commit()
            print("字段添加成功！")
        else:
            print("stream_enabled 字段已存在，无需添加。")
    else:
        print("ai_models 表不存在，请先运行 init_db.py 创建表。")
    
    # 关闭连接
    conn.close()

if __name__ == "__main__":
    add_stream_enabled_column() 